use LazyBufferCache instead of FixedSizeDiffCache #1191
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
This shouldn't be needed now? |
|
Yeah, I'll close it |
|
@ChrisRackauckas this fixes the Core2 tests. |
|
But there's a lot of other broken tests it creates? |
jClugstor
force-pushed
the
enzyme_adjoint_lbc
branch
from
4988fe1 to
a55f1de
Compare
|
I don't see any here that aren't also broken on master? |
|
What's the performance change from this? |
|
And which tests does this fix? |
|
This just fixes the Core 2 tests on Julia 1.11. As for performance I'll set a benchmark to see. |
|
I ran this: using Zygote, SciMLSensitivity
using OrdinaryDiffEq, ForwardDiff, Test
using BenchmarkTools
p = rand(3)
function dudt(u, p, t)
u .* p
end
function loss(p, sensealg)
prob = ODEProblem(dudt, [3.0, 2.0, 1.0], (0.0, 10.0), p)
sol = solve(prob, ImplicitEuler(), dt=0.01, saveat=0.1, sensealg=sensealg,
abstol=1e-5, reltol=1e-5)
sum(abs2, Array(sol))
end
@btime Zygote.gradient(
p -> loss(p, QuadratureAdjoint(autojacvec=EnzymeVJP())), p)
function lv(du, u, p, t)
du[1] = p[1] * u[1] - p[2] * u[1] * u[2]
du[2] = -3 * u[2] + u[1] * u[2]
end
function loss_lv(p,sensealg)
prob_lv = ODEProblem(lv, [1.0, 2.0], (0.0, 10.0), p)
sol = solve(prob_lv, ImplicitEuler(), dt=0.01, saveat=0.1, sensealg=sensealg,
abstol=1e-5, reltol=1e-5)
sum(abs2, Array(sol))
end
@btime Zygote.gradient(
p -> loss_lv(p, QuadratureAdjoint(autojacvec=EnzymeVJP())), [1.0, 2.0])
function vanderpol(du, u, p, t)
du[1] = u[2]
du[2] = p[1]*((1-u[1]^2)*u[2] - u[1])
end
function loss_vp(p, sensealg)
prob_vp = ODEProblem(vanderpol, [1.0, 2.0], (0.0, 10.0), p)
sol = solve(prob_vp, Rodas5(), dt=0.01, saveat=0.1, sensealg=sensealg,
abstol=1e-5, reltol=1e-5)
sum(abs2, Array(sol))
end
@btime Zygote.gradient(
p -> loss_vp(p, QuadratureAdjoint(autojacvec=EnzymeVJP())), [20.0])With FixedSizeDiffCache and LazyBufferCache on Julia 1.10. These are the results I got from |
|
Open an issue about this performance regression. It's minor enough to not be a major worry, but we should track it because if Enzyme fixes their part then we should prefer DiffCache. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Checklist
contributor guidelines, in particular the SciML Style Guide and
COLPRAC.
Additional context
This makes EnzymeVJP use LazyBufferCache instead of FixedSizeDiffCache to avoid problems with
reinterpretand Enzyme.